package main

import (
	"crypto/md5"
	"database/sql"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"

	_ "github.com/mattn/go-sqlite3"
)

// FileRecord 文件记录结构，对应数据库表结构
type FileRecord struct {
	FilePath string // 文件绝对路径
	FileMD5  string // MD5值
	FileSize int64  // 文件大小（字节）
}

type fileTask struct {
	Path  string
	Index int
	Total int
}

type Config struct {
	Directory string
	MinSize   int64
	SizeStr   string
	Check     bool
	Query     bool
	Workers   int
}

// parseSize 解析带单位的文件大小字符串
// 例如: "100k" -> 102400, "1M" -> 1048576
func parseSize(sizeStr string) (int64, error) {
	if sizeStr == "" {
		return 0, nil
	}

	// 使用正则表达式匹配数字和可选的单位后缀
	re := regexp.MustCompile(`^(\d+)([kKmMgG])?$`)
	matches := re.FindStringSubmatch(sizeStr)

	if len(matches) == 0 {
		return 0, fmt.Errorf("无效的大小格式: %s", sizeStr)
	}

	// 解析数字部分
	value, err := strconv.ParseInt(matches[1], 10, 64)
	if err != nil {
		return 0, fmt.Errorf("无法解析大小值: %s", sizeStr)
	}

	// 解析单位部分，计算实际字节数
	var unit string
	if len(matches) > 2 && matches[2] != "" {
		unit = strings.ToLower(matches[2])
	}

	switch unit {
	case "k":
		return value * 1024, nil
	case "m":
		return value * 1024 * 1024, nil
	case "g":
		return value * 1024 * 1024 * 1024, nil
	case "":
		return value, nil
	default:
		return 0, fmt.Errorf("不支持的大小单位: %s", unit)
	}
}

func main() {
	var config Config

	// 定义命令行参数，与Rust版本保持一致
	flag.StringVar(&config.Directory, "d", "", "要扫描的目录路径")
	flag.StringVar(&config.Directory, "directory", "", "要扫描的目录路径")
	flag.StringVar(&config.SizeStr, "s", "100k", "忽略小于指定大小的文件 (例如: 100k, 1M)")
	flag.StringVar(&config.SizeStr, "size", "100k", "忽略小于指定大小的文件 (例如: 100k, 1M)")
	flag.BoolVar(&config.Check, "c", false, "检查数据库中的文件是否存在")
	flag.BoolVar(&config.Check, "check", false, "检查数据库中的文件是否存在")
	flag.BoolVar(&config.Query, "q", false, "查询MD5相同的文件，按文件从大到小排序")
	flag.BoolVar(&config.Query, "query", false, "查询MD5相同的文件，按文件从大到小排序")

	// 默认工作线程数为CPU数目的1.5倍
	defaultWorkers := int(float64(runtime.NumCPU()) * 1.5)
	if defaultWorkers < 1 {
		defaultWorkers = 1
	}
	flag.IntVar(&config.Workers, "w", defaultWorkers, "指定工作线程数")
	flag.IntVar(&config.Workers, "workers", defaultWorkers, "指定工作线程数")

	flag.Parse()

	// 解析文件大小
	var err error
	config.MinSize, err = parseSize(config.SizeStr)
	if err != nil {
		log.Fatal("解析文件大小失败:", err)
	}

	// 初始化数据库
	db, err := initDatabase()
	if err != nil {
		log.Fatal("初始化数据库失败:", err)
	}
	defer db.Close()

	// 根据命令行参数执行不同的操作
	if config.Check {
		// 执行检查命令：检查数据库中的文件是否存在
		fmt.Println("开始检查数据库记录...")
		err = checkFiles(db)
		if err != nil {
			log.Fatal("检查文件失败:", err)
		}
	} else if config.Query {
		// 执行查询命令：查询重复文件
		fmt.Println("查询重复文件...")
		err = queryDuplicateFiles(db)
		if err != nil {
			log.Fatal("查询文件失败:", err)
		}
	} else if config.Directory != "" {
		// 执行扫描命令：扫描指定目录的文件
		err = scanDirectory(db, config)
		if err != nil {
			log.Fatal("扫描目录失败:", err)
		}
	} else {
		// 没有指定任何操作，显示帮助信息
		fmt.Println("错误: 必须指定一个操作")
		fmt.Println("")
		fmt.Println("使用示例:")
		fmt.Println("  filescanner -d <目录> [-s <文件大小>] [-w <线程数>]  扫描目录文件")
		fmt.Println("  filescanner -c                                     检查数据库中的文件")
		fmt.Println("  filescanner -q                                     查询重复文件")
		fmt.Println("")
		fmt.Println("参数说明:")
		fmt.Println("  -d, --directory <DIR>  指定要扫描的目录")
		fmt.Println("  -s, --size <SIZE>      忽略小于指定大小的文件 (默认: 100k)")
		fmt.Println("  -w, --workers <NUM>    指定工作线程数 (默认: CPU核心数的1.5倍)")
		fmt.Println("  -c, --check            检查数据库中的文件是否存在")
		fmt.Println("  -q, --query            查询MD5相同的文件")
		fmt.Println("")
		fmt.Println("使用 --help 查看完整帮助信息")
		os.Exit(1)
	}
}

func initDatabase() (*sql.DB, error) {
	// 获取可执行文件所在目录，与Rust版本一致
	exePath, err := os.Executable()
	if err != nil {
		return nil, fmt.Errorf("无法获取可执行文件路径: %v", err)
	}
	exeDir := filepath.Dir(exePath)
	dbPath := filepath.Join(exeDir, "filescanner.db")

	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, err
	}

	// 创建表，确保表有正确的结构
	createTableSQL := `
	CREATE TABLE IF NOT EXISTS files (
		file_path TEXT PRIMARY KEY,
		file_md5 TEXT NOT NULL,
		file_size INTEGER NOT NULL
	);`

	_, err = db.Exec(createTableSQL)
	if err != nil {
		return nil, err
	}

	// 创建索引以加速查询
	createIndexSQL := `CREATE INDEX IF NOT EXISTS idx_file_md5 ON files(file_md5);`
	_, err = db.Exec(createIndexSQL)
	if err != nil {
		return nil, err
	}

	return db, nil
}

func scanDirectory(db *sql.DB, config Config) error {
	fmt.Printf("扫描目录: \"%s\"\n", config.Directory)
	fmt.Printf("最小文件大小: %d 字节\n", config.MinSize)
	fmt.Printf("工作线程数: %d\n", config.Workers)

	// 首先收集所有符合条件的文件
	var filesToProcess []string

	err := filepath.Walk(config.Directory, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			log.Printf("访问 %s 出错: %v", path, err)
			return nil
		}

		if info.IsDir() {
			return nil
		}

		// 检查文件大小限制
		if info.Size() >= config.MinSize {
			absPath, err := filepath.Abs(path)
			if err != nil {
				log.Printf("获取绝对路径失败 %s: %v", path, err)
				return nil
			}
			filesToProcess = append(filesToProcess, absPath)
		}

		return nil
	})

	if err != nil {
		return fmt.Errorf("遍历目录失败: %v", err)
	}

	totalFiles := len(filesToProcess)
	fmt.Printf("找到符合条件的文件: %d 个\n", totalFiles)

	if totalFiles == 0 {
		return nil
	}

	// 使用通道来分发任务
	tasks := make(chan fileTask, totalFiles)
	var wg sync.WaitGroup

	// 启动工作协程
	for i := 0; i < config.Workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for task := range tasks {
				processFileTask(db, task)
			}
		}()
	}

	// 发送所有任务
	for i, filePath := range filesToProcess {
		task := fileTask{
			Path:  filePath,
			Index: i + 1,
			Total: totalFiles,
		}
		tasks <- task
	}

	close(tasks)
	wg.Wait()

	return nil
}

func processFileTask(db *sql.DB, task fileTask) {
	// 检查文件是否已在数据库中
	var existingSize int64
	query := "SELECT file_size FROM files WHERE file_path = ?"
	err := db.QueryRow(query, task.Path).Scan(&existingSize)

	if err == nil {
		// 文件已存在数据库中，检查大小是否匹配
		if stat, err := os.Stat(task.Path); err == nil && stat.Size() == existingSize {
			// 文件大小匹配，跳过处理
			fmt.Printf("[%d/%d] [Skip] %s\n", task.Index, task.Total, task.Path)
			return
		}
		// 如果文件大小不匹配，继续处理以更新记录
	}
	// 如果文件不在数据库中（err != nil），也继续处理

	// 计算文件 MD5
	fileRecord, err := processFile(task.Path)
	if err != nil {
		fmt.Printf("[%d/%d] [Skip] %s\n", task.Index, task.Total, task.Path)
		return
	}

	// 保存到数据库
	insertSQL := "INSERT OR REPLACE INTO files (file_path, file_md5, file_size) VALUES (?, ?, ?)"
	_, err = db.Exec(insertSQL, fileRecord.FilePath, fileRecord.FileMD5, fileRecord.FileSize)
	if err != nil {
		log.Printf("保存文件记录失败 %s: %v", fileRecord.FilePath, err)
		return
	}

	fmt.Printf("[%d/%d] %s\n", task.Index, task.Total, task.Path)
}

func processFile(filePath string) (FileRecord, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return FileRecord{}, err
	}
	defer file.Close()

	hash := md5.New()
	size, err := io.Copy(hash, file)
	if err != nil {
		return FileRecord{}, err
	}

	md5Hash := fmt.Sprintf("%x", hash.Sum(nil))

	return FileRecord{
		FilePath: filePath,
		FileMD5:  md5Hash,
		FileSize: size,
	}, nil
}

func checkFiles(db *sql.DB) error {
	// 获取所有文件记录
	query := "SELECT file_path, file_md5, file_size FROM files"
	rows, err := db.Query(query)
	if err != nil {
		return fmt.Errorf("获取文件记录失败: %v", err)
	}
	defer rows.Close()

	var files []FileRecord

	for rows.Next() {
		var file FileRecord
		if err := rows.Scan(&file.FilePath, &file.FileMD5, &file.FileSize); err != nil {
			log.Printf("读取文件记录失败: %v", err)
			continue
		}
		files = append(files, file)
	}

	totalFiles := len(files)
	fmt.Printf("数据库中共有 %d 个文件记录\n", totalFiles)

	if totalFiles == 0 {
		return nil
	}

	// 并发检查文件
	tasks := make(chan struct {
		FilePath string
		FileSize int64
		Index    int
	}, totalFiles)

	var wg sync.WaitGroup

	// 启动工作协程
	for i := 0; i < runtime.NumCPU(); i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for task := range tasks {
				checkFileTask(db, task.FilePath, task.FileSize, task.Index, totalFiles)
			}
		}()
	}

	// 发送任务
	for i, file := range files {
		task := struct {
			FilePath string
			FileSize int64
			Index    int
		}{
			FilePath: file.FilePath,
			FileSize: file.FileSize,
			Index:    i + 1,
		}
		tasks <- task
	}

	close(tasks)
	wg.Wait()

	return nil
}

func checkFileTask(db *sql.DB, filePath string, expectedSize int64, index, total int) {
	stat, err := os.Stat(filePath)

	if err != nil || stat.Size() != expectedSize {
		// 文件不存在或大小不匹配，删除记录
		deleteSQL := "DELETE FROM files WHERE file_path = ?"
		if _, err := db.Exec(deleteSQL, filePath); err != nil {
			log.Printf("删除文件记录失败 %s: %v", filePath, err)
		}
		fmt.Printf("[D] [%d/%d] %s\n", index, total, filePath)
	} else {
		// 文件正常
		fmt.Printf("[N] [%d/%d] %s\n", index, total, filePath)
	}
}

// queryDuplicateFiles 查询MD5相同的文件并按文件从大到小排序
func queryDuplicateFiles(db *sql.DB) error {
	// 首先查找重复的MD5值
	duplicateMD5Query := `
	SELECT file_md5 FROM files
	GROUP BY file_md5
	HAVING COUNT(*) > 1
	`

	rows, err := db.Query(duplicateMD5Query)
	if err != nil {
		return fmt.Errorf("查询重复MD5失败: %v", err)
	}
	defer rows.Close()

	var duplicateMD5s []string
	for rows.Next() {
		var md5 string
		if err := rows.Scan(&md5); err != nil {
			log.Printf("读取MD5值失败: %v", err)
			continue
		}
		duplicateMD5s = append(duplicateMD5s, md5)
	}

	if len(duplicateMD5s) == 0 {
		fmt.Println("没有找到MD5相同的文件")
		return nil
	}

	// 收集所有重复文件
	var allDuplicateFiles []FileRecord

	for _, md5 := range duplicateMD5s {
		filesQuery := `
		SELECT file_path, file_md5, file_size
		FROM files
		WHERE file_md5 = ?
		ORDER BY file_size DESC
		`

		fileRows, err := db.Query(filesQuery, md5)
		if err != nil {
			log.Printf("查询文件失败: %v", err)
			continue
		}

		for fileRows.Next() {
			var fileRecord FileRecord
			if err := fileRows.Scan(&fileRecord.FilePath, &fileRecord.FileMD5, &fileRecord.FileSize); err != nil {
				log.Printf("读取文件信息失败: %v", err)
				continue
			}
			allDuplicateFiles = append(allDuplicateFiles, fileRecord)
		}
		fileRows.Close()
	}

	// 按文件大小从大到小排序（虽然已经在每个组内排序，但需要全局排序）
	sort.Slice(allDuplicateFiles, func(i, j int) bool {
		return allDuplicateFiles[i].FileSize > allDuplicateFiles[j].FileSize
	})

	// 按照要求的格式输出：file_path|file_md5|file_size
	for _, file := range allDuplicateFiles {
		fmt.Printf("%s|%s|%d\n", file.FilePath, file.FileMD5, file.FileSize)
	}

	return nil
}
